import pandas as pd
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split, RandomizedSearchCV, GridSearchCV
import matplotlib.pyplot as plt


df = pd.read_csv('admission.csv').values
X = df[:, 1:]
y = df[:, 0]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=10)

leaves = [10, 20, 30]
params = {
    'criterion': ['entropy', 'gini'],
    'max_depth': [10, 20, 30, 40, 50],
    'max_leaf_nodes': [10, 20, 30]
}

tree = DecisionTreeClassifier()
clf = GridSearchCV(tree, params)
clf.fit(X_train, y_train)
y_predict = clf.predict(X_test)
print(accuracy_score(y_test, y_predict))

best_tree = clf.best_estimator_
print(best_tree)
feature_names = ['gre', 'gpa', 'ses', 'gander', 'race', 'rank']
plot_tree(best_tree, feature_names=feature_names)
plt.show()
